import pandas as pd
import numpy as np
from datetime import datetime
import logging

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class DataPreprocessor:
    def preprocess_daily_data(self, df):
        """
        预处理日频数据
        :param df: 原始日频数据DataFrame
        :return: 预处理后的日频数据DataFrame
        """
        if df.empty:
            logger.warning("日频数据为空，跳过预处理")
            return df
        
        try:
            # 按日期升序排序
            df = df.sort_values('trade_date')
            
            # 转换日期格式
            df['trade_date'] = pd.to_datetime(df['trade_date'], format='%Y%m%d')
            
            # 计算每日收益率
            df['daily_return'] = df['close'].pct_change()
            
            # 计算缺失值比例
            missing_ratio = df.isnull().mean()
            
            # 处理缺失值
            for col in df.columns:
                if missing_ratio[col] > 0:
                    logger.info(f"{col}列存在缺失值，缺失比例: {missing_ratio[col]:.2%}")
                    # 对于数值型数据，使用前向填充
                    if df[col].dtype in [np.float64, np.int64]:
                        df[col] = df[col].fillna(method='ffill')
                    # 对于其他类型数据，填充为0
                    else:
                        df[col] = df[col].fillna(0)
            
            # 计算技术指标
            df = self._calculate_technical_indicators(df)
            
            logger.info(f"日频数据预处理完成，共{len(df)}条记录")
            return df
        except Exception as e:
            logger.error(f"日频数据预处理失败: {e}")
            return df
    
    def _calculate_technical_indicators(self, df):
        """
        计算常用技术指标
        :param df: 日频数据DataFrame
        :return: 包含技术指标的日频数据DataFrame
        """
        # 移动平均线
        df['MA5'] = df['close'].rolling(window=5).mean()
        df['MA10'] = df['close'].rolling(window=10).mean()
        df['MA20'] = df['close'].rolling(window=20).mean()
        df['MA30'] = df['close'].rolling(window=30).mean()
        
        # MACD指标
        df['EMA12'] = df['close'].ewm(span=12, adjust=False).mean()
        df['EMA26'] = df['close'].ewm(span=26, adjust=False).mean()
        df['DIF'] = df['EMA12'] - df['EMA26']
        df['DEA'] = df['DIF'].ewm(span=9, adjust=False).mean()
        df['MACD'] = 2 * (df['DIF'] - df['DEA'])
        
        # KDJ指标
        low_9 = df['low'].rolling(window=9).min()
        high_9 = df['high'].rolling(window=9).max()
        df['RSV'] = (df['close'] - low_9) / (high_9 - low_9) * 100
        df['K'] = df['RSV'].ewm(com=2).mean()
        df['D'] = df['K'].ewm(com=2).mean()
        df['J'] = 3 * df['K'] - 2 * df['D']
        
        # RSI指标
        delta = df['close'].diff()
        gain = delta.where(delta > 0, 0)
        loss = -delta.where(delta < 0, 0)
        avg_gain = gain.rolling(window=14).mean()
        avg_loss = loss.rolling(window=14).mean()
        rs = avg_gain / avg_loss
        df['RSI'] = 100 - (100 / (1 + rs))
        
        # 成交量指标
        df['volume_MA5'] = df['vol'].rolling(window=5).mean()
        df['volume_MA10'] = df['vol'].rolling(window=10).mean()
        
        # 布林带
        df['BB_MA'] = df['close'].rolling(window=20).mean()
        df['BB_std'] = df['close'].rolling(window=20).std()
        df['BB_upper'] = df['BB_MA'] + 2 * df['BB_std']
        df['BB_lower'] = df['BB_MA'] - 2 * df['BB_std']
        
        # 移除包含NaN的行
        df = df.dropna()
        
        return df
    
    def preprocess_financial_data(self, df, daily_df):
        """
        预处理财务数据，并与日频数据合并
        :param df: 原始财务数据DataFrame
        :param daily_df: 日频数据DataFrame
        :return: 合并后的DataFrame
        """
        if df.empty or daily_df.empty:
            logger.warning("财务数据或日频数据为空，跳过预处理")
            return daily_df
        
        try:
            # 转换日期格式
            df['ann_date'] = pd.to_datetime(df['ann_date'])
            df['end_date'] = pd.to_datetime(df['end_date'])
            
            # 使用公告日期作为财务数据生效日期
            df['effective_date'] = df['ann_date']
            
            # 对于缺失公告日期的记录，使用报告期结束日期+90天作为生效日期
            mask = df['effective_date'].isnull()
            df.loc[mask, 'effective_date'] = df.loc[mask, 'end_date'] + pd.Timedelta(days=90)
            
            # 按生效日期排序
            df = df.sort_values('effective_date')
            
            # 为每个交易日匹配最近的财务数据
            merged_df = pd.merge_asof(
                daily_df, 
                df, 
                left_on='trade_date', 
                right_on='effective_date',
                direction='backward'
            )
            
            # 选择需要的财务指标
            selected_financial_cols = [
                'roe', 'roa', 'gross_margin', 'net_margin', 'debt_to_assets', 
                'current_ratio', 'quick_ratio', 'eps', 'bps', 'revenue_grow', 
                'profit_grow', 'assets_grow', 'oper_revenue', 'oper_profit', 
                'total_profit', 'net_profit', 'total_assets', 'total_liab',
                'holders_num'
            ]
            
            # 确保合并后的DataFrame包含所有需要的列
            for col in selected_financial_cols:
                if col not in merged_df.columns:
                    merged_df[col] = np.nan
            
            # 处理缺失值
            for col in selected_financial_cols:
                if merged_df[col].isnull().any():
                    merged_df[col] = merged_df[col].fillna(method='ffill')
            
            logger.info(f"财务数据预处理完成，共{len(merged_df)}条记录")
            return merged_df
        except Exception as e:
            logger.error(f"财务数据预处理失败: {e}")
            return daily_df
    
    def prepare_ml_data(self, df, target_col='daily_return', lookback_days=30, forecast_days=5):
        """
        准备机器学习模型的训练数据
        :param df: 预处理后的DataFrame
        :param target_col: 目标列
        :param lookback_days: 回望天数
        :param forecast_days: 预测天数
        :return: 特征矩阵X和目标向量y
        """
        if df.empty:
            logger.warning("数据为空，无法准备机器学习数据")
            return None, None
        
        try:
            # 选择用于建模的特征列
            feature_cols = [
                'open', 'high', 'low', 'close', 'vol', 'amount', 'MA5', 'MA10', 
                'MA20', 'MA30', 'DIF', 'DEA', 'MACD', 'K', 'D', 'J', 'RSI', 
                'volume_MA5', 'volume_MA10', 'BB_MA', 'BB_std', 'BB_upper', 'BB_lower',
                'roe', 'roa', 'gross_margin', 'net_margin', 'debt_to_assets', 
                'current_ratio', 'quick_ratio', 'eps', 'bps', 'revenue_grow', 
                'profit_grow', 'assets_grow'
            ]
            
            # 确保所有特征列都存在
            feature_cols = [col for col in feature_cols if col in df.columns]
            
            # 创建特征矩阵和目标向量
            X = []
            y = []
            
            # 遍历数据，创建时序样本
            for i in range(len(df) - lookback_days - forecast_days + 1):
                # 获取特征
                features = df.iloc[i:i+lookback_days][feature_cols].values
                X.append(features)
                
                # 获取目标
                if forecast_days == 1:
                    target = df.iloc[i+lookback_days][target_col]
                else:
                    # 多日收益率计算
                    end_price = df.iloc[i+lookback_days+forecast_days-1]['close']
                    start_price = df.iloc[i+lookback_days-1]['close']
                    target = (end_price / start_price) - 1
                
                y.append(target)
            
            # 转换为numpy数组
            X = np.array(X)
            y = np.array(y)
            
            logger.info(f"机器学习数据准备完成，样本数: {len(X)}")
            return X, y
        except Exception as e:
            logger.error(f"准备机器学习数据失败: {e}")
            return None, None    